{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "sdf_2d.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "P-ghrATIVr4W",
        "cellView": "form"
      },
      "source": [
        "# @title Install Packages.\n",
        "!pip install git+https://github.com/fogleman/sdf.git\n",
        "!pip install flax optax\n",
        "!pip install mediapy\n",
        "!pip install jax\n",
        "!pip install git+https://github.com/google/nerfies.git\n",
        "!pip install --upgrade scikit-image Pillow\n",
        "!pip install trimesh[easy]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aiQoBvSaXo5t",
        "cellView": "form"
      },
      "source": [
        "# @title Visualization utilities.\n",
        "import functools\n",
        "\n",
        "from matplotlib import cm\n",
        "from matplotlib.colors import LinearSegmentedColormap\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "_colormap_cache = {}\n",
        "\n",
        "\n",
        "def _build_colormap(name, num_bins=256):\n",
        "  base = cm.get_cmap(name)\n",
        "  color_list = base(np.linspace(0, 1, num_bins))\n",
        "  cmap_name = base.name + str(num_bins)\n",
        "  colormap = LinearSegmentedColormap.from_list(cmap_name, color_list, num_bins)\n",
        "  colormap = colormap(np.linspace(0, 1, num_bins))[:, :3]\n",
        "  return colormap\n",
        "\n",
        "\n",
        "@functools.lru_cache(maxsize=32)\n",
        "def get_colormap(name, num_bins=256):\n",
        "  \"\"\"Lazily initializes and returns a colormap.\"\"\"\n",
        "  if name == 'turbo':\n",
        "    return _TURBO_COLORS\n",
        "\n",
        "  return _build_colormap(name, num_bins)\n",
        "\n",
        "\n",
        "def interpolate_colormap(values, colormap):\n",
        "  \"\"\"Interpolates the colormap given values between 0.0 and 1.0.\"\"\"\n",
        "  a = np.floor(values * 255)\n",
        "  b = (a + 1).clip(max=255)\n",
        "  f = values * 255.0 - a\n",
        "  a = a.astype(np.uint16).clip(0, 255)\n",
        "  b = b.astype(np.uint16).clip(0, 255)\n",
        "  return colormap[a] + (colormap[b] - colormap[a]) * f[..., np.newaxis]\n",
        "\n",
        "\n",
        "def scale_values(values, vmin, vmax, eps=1e-6):\n",
        "  return (values - vmin) / max(vmax - vmin, eps)\n",
        "\n",
        "\n",
        "def colorize(\n",
        "    array, cmin=None, cmax=None, cmap='magma', eps=1e-6, invert=False):\n",
        "  \"\"\"Applies a colormap to an array.\n",
        "\n",
        "  Args:\n",
        "    array: the array to apply a colormap to.\n",
        "    cmin: the minimum value of the colormap. If None will take the min.\n",
        "    cmax: the maximum value of the colormap. If None will take the max.\n",
        "    cmap: the color mapping to use.\n",
        "    eps: a small value to prevent divide by zero.\n",
        "    invert: if True will invert the colormap.\n",
        "\n",
        "  Returns:\n",
        "    a color mapped version of array.\n",
        "  \"\"\"\n",
        "  array = np.asarray(array)\n",
        "\n",
        "  if cmin is None:\n",
        "    cmin = array.min()\n",
        "  if cmax is None:\n",
        "    cmax = array.max()\n",
        "\n",
        "  x = scale_values(array, cmin, cmax, eps)\n",
        "  colormap = get_colormap(cmap)\n",
        "  colorized = interpolate_colormap(1.0 - x if invert else x, colormap)\n",
        "  colorized[x > 1.0] = 0.0 if invert else 1.0\n",
        "  colorized[x < 0.0] = 1.0 if invert else 0.0\n",
        "\n",
        "  return colorized\n",
        "\n",
        "\n",
        "def colorize_binary_logits(array, cmap=None):\n",
        "  \"\"\"Colorizes binary logits as a segmentation map.\"\"\"\n",
        "  num_classes = array.shape[-1]\n",
        "  if cmap is None:\n",
        "    if num_classes <= 8:\n",
        "      cmap = 'Set3'\n",
        "    elif num_classes <= 10:\n",
        "      cmap = 'tab10'\n",
        "    elif num_classes <= 20:\n",
        "      cmap = 'tab20'\n",
        "    else:\n",
        "      cmap = 'gist_rainbow'\n",
        "\n",
        "  colormap = get_colormap(cmap, num_classes)\n",
        "  indices = np.argmax(array, axis=-1)\n",
        "  return np.take(colormap, indices, axis=0)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y2IwMc5eyyag",
        "cellView": "form"
      },
      "source": [
        "# @title Schedules.\n",
        "\n",
        "\"\"\"Annealing Schedules.\"\"\"\n",
        "import abc\n",
        "import collections\n",
        "import copy\n",
        "import math\n",
        "from typing import Any, Iterable, List, Tuple, Union\n",
        "\n",
        "from jax import numpy as jnp\n",
        "\n",
        "\n",
        "def from_tuple(x):\n",
        "  schedule_type, *args = x\n",
        "  return SCHEDULE_MAP[schedule_type](*args)\n",
        "\n",
        "\n",
        "def from_dict(d):\n",
        "  d = copy.copy(dict(d))\n",
        "  schedule_type = d.pop('type')\n",
        "  return SCHEDULE_MAP[schedule_type](**d)\n",
        "\n",
        "\n",
        "def from_config(schedule):\n",
        "  if isinstance(schedule, Schedule):\n",
        "    return schedule\n",
        "  if isinstance(schedule, Tuple) or isinstance(schedule, List):\n",
        "    return from_tuple(schedule)\n",
        "  if isinstance(schedule, collections.Mapping):\n",
        "    return from_dict(schedule)\n",
        "\n",
        "  raise ValueError(f'Unknown type {type(schedule)}.')\n",
        "\n",
        "\n",
        "class Schedule(abc.ABC):\n",
        "  \"\"\"An interface for generic schedules..\"\"\"\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    raise NotImplementedError\n",
        "\n",
        "  def __call__(self, step):\n",
        "    return self.get(step)\n",
        "\n",
        "\n",
        "class ConstantSchedule(Schedule):\n",
        "  \"\"\"Linearly scaled scheduler.\"\"\"\n",
        "\n",
        "  def __init__(self, value):\n",
        "    super().__init__()\n",
        "    self.value = value\n",
        "\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    return jnp.full_like(step, self.value, dtype=jnp.float32)\n",
        "\n",
        "\n",
        "class LinearSchedule(Schedule):\n",
        "  \"\"\"Linearly scaled scheduler.\"\"\"\n",
        "\n",
        "  def __init__(self, initial_value, final_value, num_steps):\n",
        "    super().__init__()\n",
        "    self.initial_value = initial_value\n",
        "    self.final_value = final_value\n",
        "    self.num_steps = num_steps\n",
        "\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    if self.num_steps == 0:\n",
        "      return jnp.full_like(step, self.final_value, dtype=jnp.float32)\n",
        "    alpha = jnp.minimum(step / self.num_steps, 1.0)\n",
        "    return (1.0 - alpha) * self.initial_value + alpha * self.final_value\n",
        "\n",
        "\n",
        "class ExponentialSchedule(Schedule):\n",
        "  \"\"\"Exponentially decaying scheduler.\"\"\"\n",
        "\n",
        "  def __init__(self, initial_value, final_value, num_steps, eps=1e-10):\n",
        "    super().__init__()\n",
        "    if initial_value <= final_value:\n",
        "      raise ValueError('Final value must be less than initial value.')\n",
        "\n",
        "    self.initial_value = initial_value\n",
        "    self.final_value = final_value\n",
        "    self.num_steps = num_steps\n",
        "    self.eps = eps\n",
        "\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    if step >= self.num_steps:\n",
        "      return jnp.full_like(step, self.final_value, dtype=jnp.float32)\n",
        "\n",
        "    final_value = max(self.final_value, self.eps)\n",
        "    base = final_value / self.initial_value\n",
        "    exponent = step / (self.num_steps - 1)\n",
        "    if step >= self.num_steps:\n",
        "      return jnp.full_like(step, self.final_value, dtype=jnp.float32)\n",
        "    return self.initial_value * base**exponent\n",
        "\n",
        "\n",
        "class CosineEasingSchedule(Schedule):\n",
        "  \"\"\"Schedule that eases slowsly using a cosine.\"\"\"\n",
        "\n",
        "  def __init__(self, initial_value, final_value, num_steps):\n",
        "    super().__init__()\n",
        "    self.initial_value = initial_value\n",
        "    self.final_value = final_value\n",
        "    self.num_steps = num_steps\n",
        "\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    alpha = jnp.minimum(step / self.num_steps, 1.0)\n",
        "    scale = self.final_value - self.initial_value\n",
        "    x = min(max(alpha, 0.0), 1.0)\n",
        "    return (self.initial_value\n",
        "            + scale * 0.5 * (1 + math.cos(jnp.pi * x + jnp.pi)))\n",
        "\n",
        "\n",
        "class StepSchedule(Schedule):\n",
        "  \"\"\"Schedule that eases slowsly using a cosine.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               initial_value,\n",
        "               decay_interval,\n",
        "               decay_factor,\n",
        "               max_decays,\n",
        "               final_value=None):\n",
        "    super().__init__()\n",
        "    self.initial_value = initial_value\n",
        "    self.decay_factor = decay_factor\n",
        "    self.decay_interval = decay_interval\n",
        "    self.max_decays = max_decays\n",
        "    if final_value is None:\n",
        "      final_value = self.initial_value * self.decay_factor**self.max_decays\n",
        "    self.final_value = final_value\n",
        "\n",
        "  def get(self, step):\n",
        "    \"\"\"Get the value for the given step.\"\"\"\n",
        "    phase = step // self.decay_interval\n",
        "    if phase >= self.max_decays:\n",
        "      return self.final_value\n",
        "    else:\n",
        "      return self.initial_value * self.decay_factor**phase\n",
        "\n",
        "\n",
        "class PiecewiseSchedule(Schedule):\n",
        "  \"\"\"A piecewise combination of multiple schedules.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self, schedules: Iterable[Tuple[int, Union[Schedule, Iterable[Any]]]]):\n",
        "    self.schedules = [from_config(s) for ms, s in schedules]\n",
        "    milestones = jnp.array([ms for ms, s in schedules])\n",
        "    self.milestones = jnp.cumsum(milestones)[:-1]\n",
        "\n",
        "  def get(self, step):\n",
        "    idx = jnp.searchsorted(self.milestones, step, side='right')\n",
        "    schedule = self.schedules[idx]\n",
        "    base_idx = self.milestones[idx - 1] if idx >= 1 else 0\n",
        "    return schedule.get(step - base_idx)\n",
        "\n",
        "\n",
        "class DelayedSchedule(Schedule):\n",
        "  \"\"\"Delays the start of the base schedule.\"\"\"\n",
        "\n",
        "  def __init__(self, base_schedule: Schedule, delay_steps, delay_mult):\n",
        "    self.base_schedule = from_config(base_schedule)\n",
        "    self.delay_steps = delay_steps\n",
        "    self.delay_mult = delay_mult\n",
        "\n",
        "  def get(self, step):\n",
        "    delay_rate = (\n",
        "        self.delay_mult\n",
        "        + (1 - self.delay_mult)\n",
        "        * jnp.sin(0.5 * jnp.pi * jnp.clip(step / self.delay_steps, 0, 1)))\n",
        "\n",
        "    return delay_rate * self.base_schedule(step)\n",
        "\n",
        "\n",
        "SCHEDULE_MAP = {\n",
        "    'constant': ConstantSchedule,\n",
        "    'linear': LinearSchedule,\n",
        "    'exponential': ExponentialSchedule,\n",
        "    'cosine_easing': CosineEasingSchedule,\n",
        "    'step': StepSchedule,\n",
        "    'piecewise': PiecewiseSchedule,\n",
        "    'delayed': DelayedSchedule,\n",
        "}\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AhFwaOY5Vye9"
      },
      "source": [
        "import trimesh\n",
        "import sdf as sdflib\n",
        "import numpy as np\n",
        "import mediapy\n",
        "from matplotlib import pyplot as plt\n",
        "import jax\n",
        "from jax import numpy as jnp\n",
        "from jax import random\n",
        "import optax\n",
        "from nerfies import utils\n",
        "\n",
        "def matmul(a, b):\n",
        "  \"\"\"jnp.matmul defaults to bfloat16, but this helper function doesn't.\"\"\"\n",
        "  return jnp.matmul(a, b, precision=jax.lax.Precision.HIGHEST)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GdLi5g7mhn2X"
      },
      "source": [
        "!wget -O OpenSans-ExtraBold.ttf \"https://github.com/opensourcedesign/fonts/blob/master/OpenSans/OpenSans-ExtraBold.ttf?raw=true\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4CkDwFs6WMxy"
      },
      "source": [
        "from numpy.core.multiarray import empty_like\n",
        "import itertools\n",
        "\n",
        "trunc = 0.05\n",
        "x, y = np.meshgrid(\n",
        "    np.linspace(-1, 1, 256),\n",
        "    np.linspace(-0.5, 0.5, 128),\n",
        ")\n",
        "coords = np.stack([x, y], axis=-1).reshape((-1, 2))\n",
        "\n",
        "# sdfs = [\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'A', 1),\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'B', 1.0),\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'C', 1.0),\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'D', 1.0),\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'E', 1.0),\n",
        "# ]\n",
        "\n",
        "# sdfs = [\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', f'{i}', 1)\n",
        "#   for i in range(0, 3)\n",
        "# ]\n",
        "\n",
        "a_fn = lambda x, y: sdflib.circle(0.3, (x, y))\n",
        "b_fn = lambda x, y: sdflib.d2.translate(sdflib.d2.rounded_x(0.3, 0.1), (x, y))\n",
        "# box_fn = lambda x, y: (sdflib.box(0.6, (x, y)) - sdflib.circle(0.1, (x, y)))\n",
        "# triangle_fn = lambda x, y: sdflib.d2.translate(sdflib.d2.scale(sdflib.equilateral_triangle(), 0.4), (x, y))\n",
        "# corners = circle_fn(-0.5, 0.5) | box_fn(0.5, 0.5) | circle_fn(0.5, -0.5) | box_fn(-0.5, -0.5)\n",
        "# sdfs = [\n",
        "#   a(-0.5, 0.0) | b(0.5, 0.0) \n",
        "#   for a, b  in itertools.product([a_fn, b_fn], [a_fn, b_fn])\n",
        "# ]\n",
        "\n",
        "# sdfs = [\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'o', 1),\n",
        "#   sdflib.text('OpenSans-ExtraBold.ttf', 'c', 1.0),\n",
        "# ]\n",
        "sdfs = [\n",
        "      sdflib.circle(0.4, (-t, 0.0)) | (sdflib.circle(0.4, (t, 0.0)) - sdflib.circle(0.1, (t, 0.0)))\n",
        "      for t in np.linspace(0.0, 0.5, 3)\n",
        "]\n",
        "# sdfs = [\n",
        "#   # sdflib.d2.elongate(sdflib.rounded_rectangle(0.7, 0.4), np.array([0.5, 0.1])),\n",
        "#   sdflib.circle(0.4, (0, 0.0)) | sdflib.circle(0.4, (-0.5, 0.0)) | sdflib.circle(0.4, (0.5, 0.0)),\n",
        "#   sdflib.circle(0.4, (-0.5, 0.0)) | sdflib.circle(0.4, (0.5, 0.0)),\n",
        "#   sdflib.circle(0.3, (-0.5, 0.0)) | sdflib.circle(0.3, (0.5, 0.0)),\n",
        "#   sdflib.circle(0.2, (-0.5, 0.0)) | sdflib.circle(0.3, (0.5, 0.0)),\n",
        "#   sdflib.circle(0.1, (-0.5, 0.0)) | sdflib.circle(0.3, (0.5, 0.0)),\n",
        "#   sdflib.circle(0.0, (0.0, 0.0)) - sdflib.circle(0.1, (0.0, 0.0)),\n",
        "# ]\n",
        "# sdfs += [\n",
        "#       sdflib.circle(0.4, (-t, 0.0)) | (sdflib.circle(0.4, (t, 0.0)))\n",
        "#       for t in np.linspace(0.0, 0.5, 3)\n",
        "# ]\n",
        "# sdfs = [\n",
        "#       sdflib.circle(0.4, (0.0, 0.0)),\n",
        "#       sdflib.circle(0.3, (-0.5, 0.0)) | sdflib.circle(0.3, (0.5, 0.0)),\n",
        "#       sdflib.circle(0.2, (-0.6, 0.0)) | sdflib.circle(0.2, (0.6, 0.0)) | sdflib.circle(0.2, (0.0, 0.0)),\n",
        "#       # sdflib.circle(0.5, (-0.0, 0.0)) - sdflib.circle(0.2, (0.0, 0.0)),\n",
        "#       sdflib.text('OpenSans-ExtraBold.ttf', 'SDF', 1.5),\n",
        "# ]\n",
        "\n",
        "sdf_images = [\n",
        "  colorize(sdf(coords).clip(-trunc, trunc).reshape(x.shape), cmin=-trunc, cmax=trunc, cmap='gist_gray', invert=False)\n",
        "  for sdf in sdfs]\n",
        "mediapy.show_images([np.flipud(x) for x in sdf_images], columns=5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pkOuY2HEWON2",
        "cellView": "form"
      },
      "source": [
        "# @title MLP\n",
        "from typing import Any, Optional, Tuple, Callable\n",
        "from flax import linen as nn\n",
        "\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  \"\"\"Basic MLP class with hidden layers and an output layers.\"\"\"\n",
        "  depth: int\n",
        "  width: int\n",
        "  hidden_init: Any = nn.initializers.xavier_uniform()\n",
        "  hidden_activation: Any = nn.relu\n",
        "  hidden_norm: Optional[Callable[[Any], nn.Module]] = None\n",
        "  output_init: Optional[Any] = None\n",
        "  output_channels: int = 0\n",
        "  output_activation: Optional[Any] = lambda x: x\n",
        "  use_bias: bool = True\n",
        "  skips: Tuple[int] = tuple()\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    inputs = x\n",
        "    for i in range(self.depth):\n",
        "      layer = nn.Dense(\n",
        "          self.width,\n",
        "          use_bias=self.use_bias,\n",
        "          kernel_init=self.hidden_init,\n",
        "          name=f'hidden_{i}')\n",
        "      if i in self.skips:\n",
        "        x = jnp.concatenate([x, inputs], axis=-1)\n",
        "      x = layer(x)\n",
        "      if self.hidden_norm:\n",
        "        x = self.hidden_norm()(x)\n",
        "      x = self.hidden_activation(x)\n",
        "\n",
        "    if self.output_channels > 0:\n",
        "      logit_layer = nn.Dense(\n",
        "          self.output_channels,\n",
        "          use_bias=self.use_bias,\n",
        "          kernel_init=self.output_init,\n",
        "          name='logit')\n",
        "      x = logit_layer(x)\n",
        "      if self.output_activation is not None:\n",
        "        x = self.output_activation(x)\n",
        "\n",
        "    return x\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fde00MEPbjVp",
        "cellView": "form"
      },
      "source": [
        "# @title Position Encoding\n",
        "\n",
        "def posenc(x, min_deg, max_deg, use_identity=False, alpha=None):\n",
        "  \"\"\"Encode `x` with sinusoids scaled by 2^[min_deg:max_deg-1].\"\"\"\n",
        "  batch_shape = x.shape[:-1]\n",
        "  scales = 2.0 ** jnp.arange(min_deg, max_deg)\n",
        "  # (*, F, C).\n",
        "  xb = x[..., None, :] * scales[:, None]\n",
        "  # (*, F, 2, C).\n",
        "  four_feat = jnp.sin(jnp.stack([xb, xb + 0.5 * jnp.pi], axis=-2))\n",
        "\n",
        "  if alpha is not None:\n",
        "    window = posenc_window(min_deg, max_deg, alpha)\n",
        "    four_feat = window[..., None, None] * four_feat\n",
        "\n",
        "  # (*, 2*F*C).\n",
        "  four_feat = four_feat.reshape((*batch_shape, -1))\n",
        "\n",
        "  if use_identity:\n",
        "    return jnp.concatenate([x, four_feat], axis=-1)\n",
        "  else:\n",
        "    return four_feat\n",
        "\n",
        "\n",
        "def posenc_window(min_deg, max_deg, alpha):\n",
        "  \"\"\"Windows a posenc using a cosiney window.\n",
        "\n",
        "  This is equivalent to taking a truncated Hann window and sliding it to the\n",
        "  right along the frequency spectrum.\n",
        "\n",
        "  Args:\n",
        "    min_deg: the lower frequency band.\n",
        "    max_deg: the upper frequency band.\n",
        "    alpha: will ease in each frequency as alpha goes from 0.0 to num_freqs.\n",
        "\n",
        "  Returns:\n",
        "    A 1-d numpy array with num_sample elements containing the window.\n",
        "  \"\"\"\n",
        "  bands = jnp.arange(min_deg, max_deg)\n",
        "  x = jnp.clip(alpha - bands, 0.0, 1.0)\n",
        "  return 0.5 * (1 + jnp.cos(jnp.pi * x + jnp.pi))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "58Tb5TQybr8f",
        "cellView": "form"
      },
      "source": [
        "# @title Model\n",
        "import jax\n",
        "\n",
        "class HyperSheetMLP(nn.Module):\n",
        "  \"\"\"An MLP that defines a bendy slicing surface through hyper space.\"\"\"\n",
        "  output_channels: int\n",
        "  min_deg: int = 0\n",
        "  max_deg: int = 1\n",
        "\n",
        "  depth: int = 4\n",
        "  width: int = 64\n",
        "  skips: Tuple[int] = (3,)\n",
        "  hidden_init: Any = jax.nn.initializers.glorot_uniform()\n",
        "  # output_init: Any = jax.nn.initializers.glorot_uniform()\n",
        "  output_init: Any = jax.nn.initializers.normal(0.05)\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, points, embed):\n",
        "    points_feat = posenc(points, self.min_deg, self.max_deg)\n",
        "    inputs = jnp.concatenate([points_feat, embed], axis=-1)\n",
        "    mlp = MLP(depth=self.depth,\n",
        "              width=self.width,\n",
        "              skips=self.skips,\n",
        "              hidden_init=self.hidden_init,\n",
        "              output_channels=self.output_channels,\n",
        "              output_init=self.output_init)\n",
        "    return mlp(inputs)\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "  num_glo_embeddings: int\n",
        "  num_glo_features: int = 8\n",
        "  spatial_min_deg: int = 0\n",
        "  spatial_max_deg: int = 8\n",
        "  hyper_min_deg: int = 0\n",
        "  hyper_max_deg: int = 3\n",
        "  hyper_num_dims: int = 1\n",
        "\n",
        "  hyper_slice_method: str = 'axis_aligned_plane'\n",
        "\n",
        "  # embedding_init: Any = jax.nn.initializers.uniform(scale=0.05)\n",
        "  embedding_init: Any = jax.nn.initializers.glorot_uniform()\n",
        "\n",
        "  def setup(self):\n",
        "    self.template_mlp = MLP(\n",
        "      depth=8, \n",
        "      width=128, \n",
        "      hidden_init=jax.nn.initializers.glorot_uniform(),\n",
        "      # hidden_norm=nn.LayerNorm,\n",
        "      output_init=jax.nn.initializers.glorot_uniform(),\n",
        "      # output_activation=nn.tanh,\n",
        "      skips=(4,),\n",
        "      output_channels=1)\n",
        "    self.hyper_embed = nn.Embed(\n",
        "        num_embeddings=self.num_glo_embeddings,\n",
        "        features=(self.hyper_num_dims \n",
        "                  if self.hyper_slice_method == 'axis_aligned_plane' \n",
        "                  else self.num_glo_features),\n",
        "        embedding_init=self.embedding_init)\n",
        "    self.hyper_sheet_mlp = HyperSheetMLP(output_channels=self.hyper_num_dims)\n",
        "    \n",
        "  def eval_template(self, x):\n",
        "    x, z = jnp.split(x, (2,), axis=-1)\n",
        "    x = jnp.concatenate([\n",
        "        posenc(x, self.spatial_min_deg, self.spatial_max_deg), \n",
        "        posenc(z, self.hyper_min_deg, self.hyper_max_deg),\n",
        "    ], axis=-1)\n",
        "    return self.template_mlp(x)\n",
        "\n",
        "  def compute_latent(self, glo_id):\n",
        "    return self.hyper_encoder(glo_id)\n",
        "\n",
        "  def evaluate(self, x, z):\n",
        "    x = jnp.concatenate([\n",
        "        posenc(x, self.spatial_min_deg, self.spatial_max_deg), \n",
        "        posenc(z, self.hyper_min_deg, self.hyper_max_deg),\n",
        "    ], axis=-1)\n",
        "    return self.template_mlp(x)\n",
        "\n",
        "  def __call__(self, x, glo_id, alpha):\n",
        "    hyper_embed = self.hyper_embed(glo_id)\n",
        "    if self.hyper_slice_method == 'axis_aligned_plane':\n",
        "      z = hyper_embed\n",
        "    elif self.hyper_slice_method == 'bendy_sheet':\n",
        "      z = self.hyper_sheet_mlp(x, hyper_embed)\n",
        "    else:\n",
        "      raise RuntimeError('Unknown hyper slice method.')\n",
        "    \n",
        "    return {\n",
        "        'values': self.evaluate(x, z),\n",
        "        'z': z,\n",
        "    }\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xjp4FS8iz6c5",
        "cellView": "form"
      },
      "source": [
        "# @title Barron loss.\n",
        "\n",
        "@jax.jit\n",
        "def general_loss(x, alpha, scale):\n",
        "  r\"\"\"Implements the general form of the loss.\n",
        "  This implements the rho(x, \\alpha, c) function described in \"A General and\n",
        "  Adaptive Robust Loss Function\", Jonathan T. Barron,\n",
        "  https://arxiv.org/abs/1701.03077.\n",
        "  Args:\n",
        "    x: The residual for which the loss is being computed. x can have any shape,\n",
        "      and alpha and scale will be broadcasted to match x's shape if necessary.\n",
        "    alpha: The shape parameter of the loss (\\alpha in the paper), where more\n",
        "      negative values produce a loss with more robust behavior (outliers \"cost\"\n",
        "      less), and more positive values produce a loss with less robust behavior\n",
        "      (outliers are penalized more heavily). Alpha can be any value in\n",
        "      [-infinity, infinity], but the gradient of the loss with respect to alpha\n",
        "      is 0 at -infinity, infinity, 0, and 2. Varying alpha allows for smooth\n",
        "      interpolation between several discrete robust losses:\n",
        "        alpha=-Infinity: Welsch/Leclerc Loss.\n",
        "        alpha=-2: Geman-McClure loss.\n",
        "        alpha=0: Cauchy/Lortentzian loss.\n",
        "        alpha=1: Charbonnier/pseudo-Huber loss.\n",
        "        alpha=2: L2 loss.\n",
        "    scale: The scale parameter of the loss. When |x| < scale, the loss is an\n",
        "      L2-like quadratic bowl, and when |x| > scale the loss function takes on a\n",
        "      different shape according to alpha.\n",
        "  Returns:\n",
        "    The losses for each element of x, in the same shape as x.\n",
        "  \"\"\"\n",
        "  eps = jnp.finfo(jnp.float32).eps\n",
        "\n",
        "  # `scale` must be > 0.\n",
        "  scale = jnp.maximum(eps, scale)\n",
        "\n",
        "  # The loss when alpha == 2. This will get reused repeatedly.\n",
        "  loss_two = 0.5 * (x / scale)**2\n",
        "\n",
        "  # \"Safe\" versions of log1p and expm1 that will not NaN-out.\n",
        "  log1p_safe = lambda x: jnp.log1p(jnp.minimum(x, 3e37))\n",
        "  expm1_safe = lambda x: jnp.expm1(jnp.minimum(x, 87.5))\n",
        "\n",
        "  # The loss when not in one of the special casess.\n",
        "  # Clamp |alpha| to be >= machine epsilon so that it's safe to divide by.\n",
        "  a = jnp.where(alpha >= 0, jnp.ones_like(alpha),\n",
        "                -jnp.ones_like(alpha)) * jnp.maximum(eps, jnp.abs(alpha))\n",
        "  # Clamp |2-alpha| to be >= machine epsilon so that it's safe to divide by.\n",
        "  b = jnp.maximum(eps, jnp.abs(alpha - 2))\n",
        "  loss_ow = (b / a) * ((loss_two / (0.5 * b) + 1)**(0.5 * alpha) - 1)\n",
        "\n",
        "  # Select which of the cases of the loss to return as a function of alpha.\n",
        "  return scale * jnp.where(\n",
        "      alpha == -jnp.inf, -expm1_safe(-loss_two),\n",
        "      jnp.where(\n",
        "          alpha == 0, log1p_safe(loss_two),\n",
        "          jnp.where(alpha == 2, loss_two,\n",
        "                    jnp.where(alpha == jnp.inf, expm1_safe(loss_two),\n",
        "                              loss_ow))))\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def general_loss_sq(x_sq, alpha, scale):\n",
        "  r\"\"\"Implements the general form of the loss.\n",
        "  This implements the rho(x, \\alpha, c) function described in \"A General and\n",
        "  Adaptive Robust Loss Function\", Jonathan T. Barron,\n",
        "  https://arxiv.org/abs/1701.03077.\n",
        "  Args:\n",
        "    x: The residual for which the loss is being computed. x can have any shape,\n",
        "      and alpha and scale will be broadcasted to match x's shape if necessary.\n",
        "    alpha: The shape parameter of the loss (\\alpha in the paper), where more\n",
        "      negative values produce a loss with more robust behavior (outliers \"cost\"\n",
        "      less), and more positive values produce a loss with less robust behavior\n",
        "      (outliers are penalized more heavily). Alpha can be any value in\n",
        "      [-infinity, infinity], but the gradient of the loss with respect to alpha\n",
        "      is 0 at -infinity, infinity, 0, and 2. Varying alpha allows for smooth\n",
        "      interpolation between several discrete robust losses:\n",
        "        alpha=-Infinity: Welsch/Leclerc Loss.\n",
        "        alpha=-2: Geman-McClure loss.\n",
        "        alpha=0: Cauchy/Lortentzian loss.\n",
        "        alpha=1: Charbonnier/pseudo-Huber loss.\n",
        "        alpha=2: L2 loss.\n",
        "    scale: The scale parameter of the loss. When |x| < scale, the loss is an\n",
        "      L2-like quadratic bowl, and when |x| > scale the loss function takes on a\n",
        "      different shape according to alpha.\n",
        "  Returns:\n",
        "    The losses for each element of x, in the same shape as x.\n",
        "  \"\"\"\n",
        "  eps = jnp.finfo(jnp.float32).eps\n",
        "\n",
        "  # `scale` must be > 0.\n",
        "  scale = jnp.maximum(eps, scale)\n",
        "\n",
        "  # The loss when alpha == 2. This will get reused repeatedly.\n",
        "  loss_two = 0.5 * x_sq / (scale ** 2)\n",
        "\n",
        "  # \"Safe\" versions of log1p and expm1 that will not NaN-out.\n",
        "  log1p_safe = lambda x: jnp.log1p(jnp.minimum(x, 3e37))\n",
        "  expm1_safe = lambda x: jnp.expm1(jnp.minimum(x, 87.5))\n",
        "\n",
        "  # The loss when not in one of the special casess.\n",
        "  # Clamp |alpha| to be >= machine epsilon so that it's safe to divide by.\n",
        "  a = jnp.where(alpha >= 0, jnp.ones_like(alpha),\n",
        "                -jnp.ones_like(alpha)) * jnp.maximum(eps, jnp.abs(alpha))\n",
        "  # Clamp |2-alpha| to be >= machine epsilon so that it's safe to divide by.\n",
        "  b = jnp.maximum(eps, jnp.abs(alpha - 2))\n",
        "  loss_ow = (b / a) * ((loss_two / (0.5 * b) + 1)**(0.5 * alpha) - 1)\n",
        "\n",
        "  # Select which of the cases of the loss to return as a function of alpha.\n",
        "  return scale * jnp.where(\n",
        "      alpha == -jnp.inf, -expm1_safe(-loss_two),\n",
        "      jnp.where(\n",
        "          alpha == 0, log1p_safe(loss_two),\n",
        "          jnp.where(alpha == 2, loss_two,\n",
        "                    jnp.where(alpha == jnp.inf, expm1_safe(loss_two),\n",
        "                              loss_ow))))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E8Wrv7qnuwlg"
      },
      "source": [
        "# Test out losses.\n",
        "x = jnp.linspace(-0.05, 0.05, 1000)\n",
        "loss_fn = lambda x: optax.huber_loss(x, delta=0.01) / 0.01\n",
        "# loss_fn = optax.l2_loss\n",
        "# loss_fn = jnp.abs\n",
        "# loss_fn = lambda x: general_loss(x, alpha=0.5, scale=0.005)\n",
        "loss_fn = lambda x: general_loss_sq(x ** 2, alpha=0.5, scale=0.005)\n",
        "y = jax.vmap(jax.grad(loss_fn))(x)\n",
        "\n",
        "plt.plot(x, y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4qSw9XZjcx61"
      },
      "source": [
        "# @title Define and initialize model.\n",
        "spatial_min_deg = 0  # @param {type: 'number'}\n",
        "spatial_max_deg = 4.0  # @param {type: 'number'}\n",
        "hyper_min_deg = 0  # @param {type: 'number'}\n",
        "hyper_max_deg = 1.0  # @param {type: 'number'}\n",
        "hyper_num_dims =   1# @param {type: 'number'}\n",
        "hyper_slice_method = 'axis_aligned_plane'  # @param ['axis_aligned_plane', 'bendy_sheet']\n",
        "\n",
        "devices = jax.devices()\n",
        "rng = random.PRNGKey(0)\n",
        "rng, key = random.split(rng, 2)\n",
        "\n",
        "model = Model(\n",
        "    num_glo_embeddings=len(sdfs),\n",
        "    spatial_min_deg=spatial_min_deg,\n",
        "    spatial_max_deg=spatial_max_deg,\n",
        "    hyper_min_deg=hyper_min_deg,\n",
        "    hyper_max_deg=hyper_max_deg,\n",
        "    hyper_num_dims=hyper_num_dims,\n",
        "    hyper_slice_method=hyper_slice_method,\n",
        ")\n",
        "\n",
        "init_coords = random.normal(key, (2,))\n",
        "init_glo_ids = jnp.zeros((), jnp.uint32)\n",
        "init_params = model.init(key, init_coords, init_glo_ids, 0.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "10f_5vqAc-Nr"
      },
      "source": [
        "# @title Setup dataset\n",
        "def prepare_data(xs):\n",
        "  \"\"\"Convert a input batch from tf Tensors to numpy arrays.\"\"\"\n",
        "  local_device_count = jax.local_device_count()\n",
        "  def _prepare(x):\n",
        "    # reshape (host_batch_size, height, width, 3) to\n",
        "    # (local_devices, device_batch_size, height, width, 3)\n",
        "    x = x.reshape((local_device_count, -1) + x.shape[1:])\n",
        "    return jax.api.device_put_sharded(list(x), jax.local_devices())\n",
        "\n",
        "  return jax.tree_map(_prepare, xs)\n",
        "\n",
        "\n",
        "def get_coords(height, width):\n",
        "  hr = height / max(height, width)\n",
        "  wr = width / max(height, width)\n",
        "  x, y = jnp.meshgrid(\n",
        "      np.linspace(-wr, wr, width),\n",
        "      np.linspace(-hr, hr, height))\n",
        "  coords = jnp.stack([x, y], axis=-1)\n",
        "  return coords\n",
        "\n",
        "\n",
        "def make_batch(key, batch_size, minval=-1.0, maxval=1.0, trunc=0.1):\n",
        "  key1, key2 = random.split(key, 2)\n",
        "  coords = random.uniform(key1, (batch_size, 2), minval=minval, maxval=maxval)\n",
        "  coords = coords.reshape((-1, 2))\n",
        "  ids = jnp.concatenate([\n",
        "      jnp.full((*coords.shape[:-1], 1), fill_value=i, dtype=jnp.uint32)\n",
        "      for i in range(len(sdfs))\n",
        "  ])\n",
        "  values = jnp.concatenate(\n",
        "      [jnp.array(sdf(np.array(coords)).clip(-trunc, trunc)) for sdf in sdfs],\n",
        "      axis=0\n",
        "  )\n",
        "  coords = jnp.tile(coords, (len(sdfs), 1))\n",
        "  perm = random.permutation(key2, np.arange(ids.shape[0]))\n",
        "  return {\n",
        "      'values': values[perm],\n",
        "      'ids': ids[perm],\n",
        "      'coords': coords[perm],\n",
        "  }\n",
        "\n",
        "\n",
        "batch = make_batch(random.PRNGKey(0), 1024)\n",
        "jax.tree_map(lambda x: x.shape, batch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gbcmzsW0fO-B"
      },
      "source": [
        "# @title Losses and training step.\n",
        "\n",
        "\n",
        "def compute_normal(fn, coords, eps=1e-15):\n",
        "  grad = jax.grad(fn)(coords)\n",
        "  norm = utils.safe_norm(grad, axis=-1)\n",
        "  return grad / jnp.maximum(eps, norm)\n",
        "\n",
        "\n",
        "def compute_curvature_loss(params, batch, key, pred, scale=0.05, num_z_samples=1):\n",
        "  key1, key2 = random.split(key, 2)\n",
        "\n",
        "  z_vals = pred['z']\n",
        "  xy_coords = batch['coords']\n",
        "  z_shape = (num_z_samples, *xy_coords.shape[:-1], hyper_num_dims)\n",
        "\n",
        "  z_coords = random.uniform(key, z_shape)\n",
        "  z_min = z_vals.min(axis=0, keepdims=True)[:, None, :]\n",
        "  z_max = z_vals.max(axis=0, keepdims=True)[:, None, :]\n",
        "  z_coords = z_coords * (z_max - z_min) + z_min\n",
        "  xy_coords = jnp.tile(xy_coords[None, ...], (z_coords.shape[0], 1, 1))\n",
        "\n",
        "  coords = jnp.concatenate([xy_coords, z_coords], axis=-1)\n",
        "  coords = coords.reshape((-1, 2 + z_shape[-1]))\n",
        "\n",
        "  def eval_template(x):\n",
        "    dummy_code = jnp.zeros((coords.shape[0],))\n",
        "    return model.apply(params, \n",
        "                       x, \n",
        "                       method=model.eval_template)[..., 0]\n",
        "  \n",
        "  normal_fn = jax.vmap(compute_normal, in_axes=(None, 0))\n",
        "  normals = normal_fn(eval_template, coords)\n",
        "  # Projection onto the tangent plane.\n",
        "  P = jnp.eye(normals.shape[-1])[None, ...] - normals[:, :, None] @ normals[:, None, :]\n",
        "\n",
        "  jitter_dir = random.normal(key2, coords.shape)\n",
        "  jitter_dir = jitter_dir.at[..., :2].set(0.0)\n",
        "  jitter_dir = matmul(P, jitter_dir[:, :, None]).squeeze(-1)\n",
        "  jitter_dir = jitter_dir / jnp.linalg.norm(jitter_dir, axis=-1, keepdims=True)\n",
        "  jittered_coords = coords + scale * jitter_dir\n",
        "  jittered_normals = normal_fn(eval_template, jittered_coords)\n",
        "  curvature = (jittered_normals - normals) / scale\n",
        "\n",
        "  # Only apply the curvature loss near the surface.\n",
        "  template_vals = jax.lax.stop_gradient(jax.vmap(eval_template)(coords))\n",
        "  curvature_weights = trunc - jnp.abs(template_vals.clip(a_min=-trunc, a_max=trunc))\n",
        "  sq_residual = jnp.sum(curvature ** 2, axis=-1)\n",
        "  # curvature_loss = general_loss(residual, alpha=1.0, scale=0.005)\n",
        "  curvature_loss = general_loss_sq(sq_residual, alpha=1.0, scale=0.005)\n",
        "  curvature_loss = curvature_weights * curvature_loss\n",
        "\n",
        "  return curvature_loss.mean(), jnp.sqrt(sq_residual)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def loss_fn(params, batch, scalar_params, key):\n",
        "  alpha = scalar_params['alpha']\n",
        "\n",
        "  def compute_sdf(coords, ids):\n",
        "    return model.apply(params, coords, ids, alpha)\n",
        "\n",
        "  pred = jax.vmap(compute_sdf)(batch['coords'], batch['ids'].squeeze(-1))\n",
        "  values = pred['values']\n",
        "  sdf_loss = jnp.abs(values - batch['values'])\n",
        "  sdf_loss = sdf_loss.mean()\n",
        "\n",
        "  hyper_reg_loss = (pred['z'] ** 2).sum(axis=-1).mean()\n",
        "\n",
        "  # curvature_loss, curvature = compute_curvature_loss(params, batch, key, pred)\n",
        "  # total_loss = (sdf_loss \n",
        "  #               + scalar_params['curvature_weight'] * curvature_loss\n",
        "  #               + 1e-5 * hyper_reg_loss)\n",
        "  total_loss = sdf_loss\n",
        "\n",
        "  stats_dict = {\n",
        "      'sdf_loss': sdf_loss, \n",
        "      # 'curvature_weight': (curvature_weight.min(), curvature_weight.max()),\n",
        "      # 'curvature_stats': (curvature.min(), curvature.mean(), curvature.max()),\n",
        "      # 'curvature_loss': curvature_loss,\n",
        "      'hyper_reg_loss': hyper_reg_loss,\n",
        "      'total_loss': total_loss,\n",
        "  }\n",
        "\n",
        "  return total_loss, stats_dict\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def train_step(optimizer, batch, key, scalar_params):\n",
        "  lr = scalar_params['lr']\n",
        "  _, key = random.split(rng, 2)\n",
        "  grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True)\n",
        "  (_, stats_dict), grad = grad_fn(optimizer.target, batch, scalar_params, key)\n",
        "  stats_dict = jax.lax.pmean(stats_dict, axis_name='batch')\n",
        "  grad = jax.lax.pmean(grad, axis_name='batch')\n",
        "  new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)\n",
        "  return new_optimizer, key, stats_dict\n",
        "\n",
        "\n",
        "def plot_sdf(model, params, alpha, columns=5):\n",
        "  def render(coords, glo_id):\n",
        "    return model.apply(params, coords, glo_id, alpha)\n",
        "\n",
        "  coords = get_coords(128, 256)\n",
        "\n",
        "  sdf_images = []\n",
        "  for glo_id in range(len(sdfs)):\n",
        "    glo_id = jnp.full(coords.shape[:-1], fill_value=glo_id, dtype=jnp.uint32)\n",
        "    values = jax.vmap(render)(coords, glo_id)['values']\n",
        "    values = values.clip(-trunc, trunc)\n",
        "    sdf_image = colorize(values.reshape(coords.shape[:2]), cmap='coolwarm', cmin=-trunc, cmax=trunc)\n",
        "    sdf_images.append(sdf_image)\n",
        "  mediapy.show_images([np.flipud(x) for x in sdf_images], columns=columns)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SZbVL_qegiN7",
        "cellView": "form"
      },
      "source": [
        "# @title Train.\n",
        "from nerfies import schedules\n",
        "from nerfies import utils\n",
        "from flax import optim\n",
        "from flax import jax_utils\n",
        "\n",
        "\n",
        "max_iters = 4000\n",
        "batch_size = 512\n",
        "lr_schedule = schedules.ConstantSchedule(1e-3)\n",
        "# lr_schedule = from_config({\n",
        "#   'type': 'delayed',\n",
        "#   'delay_steps': 50,\n",
        "#   'delay_mult': 0.01,\n",
        "#   'base_schedule': {\n",
        "#     'type': 'exponential',\n",
        "#     'initial_value': 1e-3,\n",
        "#     'final_value': 5e-4,\n",
        "#     'num_steps': max_iters,\n",
        "#   },\n",
        "# })\n",
        "\n",
        "curvature_schedule = schedules.from_config({\n",
        "    'type': 'piecewise',\n",
        "    'schedules': [\n",
        "      (100, ('constant', 0.1)),\n",
        "      (0, ('cosine_easing', 0.1, 1.0, 1000)),\n",
        "    ]\n",
        "})\n",
        "\n",
        "# curvature_schedule = schedules.LinearSchedule(\n",
        "    # 0.0, 100.0, 1000)\n",
        "\n",
        "alpha_schedule = schedules.ConstantSchedule(0.0)\n",
        "\n",
        "optimizer_def = optim.Adam(lr_schedule(0))\n",
        "optimizer_def = optim.WeightNorm(optimizer_def)\n",
        "optimizer = optimizer_def.create(init_params)\n",
        "optimizer = jax_utils.replicate(optimizer, devices)\n",
        "\n",
        "p_train_step = jax.pmap(\n",
        "    train_step, axis_name='batch', devices=devices, in_axes=(0, 0, 0, None))\n",
        "tt = utils.TimeTracker()\n",
        "\n",
        "rng = random.PRNGKey(1)\n",
        "keys = random.split(rng, len(devices))\n",
        "for i in range(max_iters):\n",
        "  rng, key = random.split(rng, 2)\n",
        "  batch = make_batch(key, batch_size)\n",
        "  batch = prepare_data(batch)\n",
        "\n",
        "  if i > max_iters:\n",
        "    break\n",
        "  scalar_params = {\n",
        "      'lr': lr_schedule(i),\n",
        "      'alpha': alpha_schedule(i),\n",
        "      'curvature_weight': curvature_schedule(i),\n",
        "  }\n",
        "  with tt.record_time('p_train_step'):\n",
        "    optimizer, keys, stats = p_train_step(optimizer, batch, keys, scalar_params)\n",
        "\n",
        "  if i % 10 == 0:\n",
        "    stats = jax_utils.unreplicate(stats)\n",
        "    stats = jax.tree_map(lambda x: x.item(), stats)\n",
        "    scalar_params = jax.tree_map(lambda x: (x if isinstance(x, float) else x.item()), scalar_params)\n",
        "    print(f'{i} scalar_params: {scalar_params}')\n",
        "    print(f'{i} stats: {stats}')\n",
        "    print('')\n",
        "\n",
        "  if i % 100 == 0:\n",
        "    plot_sdf(model, jax_utils.unreplicate(optimizer.target), scalar_params['alpha'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LD7P_mGKe-G8"
      },
      "source": [
        "plot_sdf(model, jax_utils.unreplicate(optimizer.target), scalar_params['alpha'], columns=4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7-WnwjS9MqD3"
      },
      "source": [
        "# Compute the minimum and maximum Z-axis coordinates.\n",
        "\n",
        "def compute_bounds(model, params, alpha):\n",
        "  coords = get_coords(100, 100).reshape((-1, 2))\n",
        "\n",
        "  sdf_images = []\n",
        "  zmin = jnp.array([float('inf')] * hyper_num_dims)\n",
        "  zmax = jnp.array([-float('inf')] * hyper_num_dims)\n",
        "  for glo_id in range(len(sdfs)):\n",
        "    glo_id = jnp.full(coords.shape[:-1], fill_value=glo_id, dtype=jnp.uint32)\n",
        "    z = jax.vmap(model.apply, in_axes=(None, 0, 0, None))(params, coords, glo_id, alpha)['z']\n",
        "    zmin = jnp.minimum(z.min(axis=0) , zmin)\n",
        "    zmax = jnp.maximum(z.max(axis=0) , zmax)\n",
        "  return zmin, zmax\n",
        "\n",
        "\n",
        "zmin, zmax = compute_bounds(model, jax_utils.unreplicate(optimizer.target), scalar_params['alpha'])\n",
        "zmin, zmax"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IN1qtkXsIgm3"
      },
      "source": [
        "import trimesh\n",
        "\n",
        "@sdflib.sdf3\n",
        "def template_sdf(params):\n",
        "    def _f(p):\n",
        "      if hyper_num_dims > 1:\n",
        "        p = jnp.concatenate([p, jnp.full_like(p[..., :1], fill_value=zmin[1])], axis=-1)\n",
        "      return model.apply(params,\n",
        "                        jnp.array(p), \n",
        "                        method=model.eval_template)\n",
        "\n",
        "    def f(p):\n",
        "        values = jax.vmap(_f)(p)\n",
        "        return np.array(values).clip(-trunc, trunc)\n",
        "    return f\n",
        "\n",
        "\n",
        "bounds = [(-1, -1, zmin[0]-0.01), (1, 1, zmax[0]+0.01)]\n",
        "# bounds = [(-1, -1, -1), (1, 1, 1)]\n",
        "sdf = template_sdf(jax_utils.unreplicate(optimizer.target))\n",
        "\n",
        "out_name = 'sdf.stl'\n",
        "sdf.save(out_name, bounds=bounds, samples=2**16)\n",
        "\n",
        "mesh = trimesh.load('sdf.stl')\n",
        "mesh.show(smooth=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HNNuxzs4FCC1"
      },
      "source": [
        "from nerfies import visualization as viz\n",
        "\n",
        "\n",
        "def compute_zmaps(model, params, alpha, shape=(128, 256)):\n",
        "  coords = get_coords(*shape).reshape((-1, 2))\n",
        "  zmaps = []\n",
        "  for glo_id in range(len(sdfs)):\n",
        "    glo_id = jnp.full(coords.shape[:-1], fill_value=glo_id, dtype=jnp.uint32)\n",
        "    zmap = jax.vmap(model.apply, in_axes=(None, 0, 0, None))(params, coords, glo_id, alpha)['z'].reshape(shape)\n",
        "    zmaps.append(zmap)\n",
        "  \n",
        "  return zmaps\n",
        "\n",
        "zmaps = compute_zmaps(model, jax_utils.unreplicate(optimizer.target), scalar_params['alpha'])\n",
        "mediapy.show_images([viz.colorize(np.array(zmap), cmap='magma', cmin=zmin.item(), cmax=zmax.item()) for zmap in zmaps], columns=4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EbcYES2Bytf9"
      },
      "source": [
        "import meshio\n",
        "colors = (viz.get_colormap('Pastel1', len(zmaps)) * 255).astype(np.uint8)\n",
        "\n",
        "for i, zmap in enumerate(zmaps):\n",
        "  color = colors[i]\n",
        "  vert_inds = np.arange(np.prod(zmap.shape)).reshape(zmap.shape)\n",
        "  quads = np.stack([\n",
        "      vert_inds[:-1, :-1],\n",
        "      vert_inds[:-1, 1:],\n",
        "      vert_inds[1:, 1:],\n",
        "      vert_inds[1:, :-1],\n",
        "  ], axis=-1).reshape((-1, 4))\n",
        "  points = np.concatenate([get_coords(*zmap.shape), zmap[..., None]], axis=-1).reshape((-1, 3))\n",
        "  cells = [('quad', quads.tolist())]\n",
        "\n",
        "  meshio.write_points_cells(\n",
        "      f\"plane_{i}.ply\", \n",
        "      points.tolist(), \n",
        "      cells,\n",
        "      point_data={\n",
        "          'red': np.full_like(points[..., 0], fill_value=color[0], dtype=np.uint8),\n",
        "          'green': np.full_like(points[..., 0], fill_value=color[1], dtype=np.uint8),\n",
        "          'blue': np.full_like(points[..., 0], fill_value=color[2], dtype=np.uint8),\n",
        "      })"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vo5qiAxq3SsK"
      },
      "source": [
        "filenames = [f'plane_{i}.ply' for i in range(len(zmaps))]\n",
        "filenames = ' '.join(filenames)\n",
        "!zip -m sdf.zip $filenames sdf.stl\n",
        "files.download('sdf.zip')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYges0ydy4-7"
      },
      "source": [
        "for i in range(len(zmaps)):\n",
        "  files.download(f'plane_{i}.ply')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PTbT_t-JFeMn"
      },
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "\n",
        "def eval_template(x):\n",
        "  dummy_code = jnp.zeros((x.shape[0],))\n",
        "  return model.apply(\n",
        "      jax_utils.unreplicate(optimizer.target), x, method=model.eval_template)[..., 0]\n",
        "\n",
        "\n",
        "num_z = 200\n",
        "xy = get_coords(256, 256)\n",
        "zvals = np.linspace(zmin, zmax, num_z)\n",
        "results = []\n",
        "for i in tqdm(range(num_z)):\n",
        "  z = jnp.full_like(xy[..., :1], fill_value=zvals[i])\n",
        "  coords = jnp.concatenate([xy, z], axis=-1).reshape((-1, 3))\n",
        "  values = jax.vmap(eval_template)(coords)\n",
        "  values = values.reshape(xy.shape[:-1])\n",
        "  results.append(np.array(values))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PGU-7HRzXQuV"
      },
      "source": [
        "from nerfies import visualization as viz\n",
        "mediapy.show_video([viz.colorize(p.clip(-trunc, trunc), cmin=-trunc, cmax=trunc, cmap='coolwarm') for p in results], codec='gif', fps=30)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JbyM_RbwhL4v"
      },
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "\n",
        "def eval_template(x):\n",
        "  dummy_code = jnp.zeros((x.shape[0],))\n",
        "  return model.apply(\n",
        "      jax_utils.unreplicate(optimizer.target), x, method=model.eval_template)[..., 0]\n",
        "\n",
        "\n",
        "num_z = 3\n",
        "xy = get_coords(64, 64)\n",
        "zgrid = jnp.meshgrid(*[\n",
        "  np.linspace(zmin[i], zmax[i], num_z) for i in range(hyper_num_dims)\n",
        "])\n",
        "zvals = jnp.stack(zgrid, axis=-1).reshape((-1, hyper_num_dims))\n",
        "results = []\n",
        "for zval in tqdm(zvals):\n",
        "  z = jnp.tile(zval[None, None, :], (xy.shape[0], xy.shape[1], 1))\n",
        "  coords = jnp.concatenate([xy, z], axis=-1).reshape((-1, 2 + z.shape[-1]))\n",
        "  values = jax.vmap(eval_template)(coords)\n",
        "  values = values.reshape(xy.shape[:-1])\n",
        "  results.append(np.array(values))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v_vUwJfIAqc8"
      },
      "source": [
        "from nerfies import visualization as viz\n",
        "mediapy.show_images([viz.colorize(np.flipud(p).clip(-trunc, trunc), cmin=-trunc, cmax=trunc, cmap='coolwarm') for p in results], columns=num_z)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cm8k9OLbAwOY"
      },
      "source": [
        "from scipy.interpolate import interpolate\n",
        "\n",
        "def compute_latent(glo_id):\n",
        "  return model.apply(jax_utils.unreplicate(optimizer.target), jnp.array([glo_id]), method=model.compute_latent)\n",
        "\n",
        "latent1 = compute_latent(8)\n",
        "latent2 = compute_latent(9)\n",
        "latent_codes = utils.interpolate_codes([latent1, latent2], 100, method='linear').squeeze()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YBghjefNqyl7"
      },
      "source": [
        "def render_code(x, z):\n",
        "  return model.apply(jax_utils.unreplicate(optimizer.target), x, z, method=model.evaluate)\n",
        "\n",
        "results = []\n",
        "for latent_code in tqdm(latent_codes):\n",
        "  x = get_coords(256, 256)\n",
        "  z = jnp.broadcast_to(latent_code[None, None, :], (*x.shape[:-1], latent_codes.shape[-1]))\n",
        "  result = jax.vmap(jax.vmap(render_code))(x, z)\n",
        "  results.append(result)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aT8Mb8WprIv6"
      },
      "source": [
        "mediapy.show_video([viz.colorize(np.flipud(p.squeeze()).clip(-trunc, trunc), cmin=-trunc, cmax=trunc, cmap='coolwarm') for p in results], codec='gif', fps=24)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ICDecjSBrJHn"
      },
      "source": [
        "results[0].shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S6F8_92Nr3RP"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}